
package com.offer.leetcode;


import java.util.Arrays;
import java.util.Collections;
import java.util.PriorityQueue;

/**
 * Solution40
 */
public class Solution40 {

    public static void main(final String[] args) {
        final Solution40 s = new Solution40();

        final int[] nums = {0, 0, 1, 2, 4, 2, 2, 3, 1, 4};

        // final int[] result = s.getLeastNumbers(nums, 8);
        final Integer[] result = s.getTopK(nums, 8);

        Arrays.stream(result).forEach(System.out::println);
    }

    public int[] getLeastNumbers(final int[] nums, final int k) {

        if (nums == null) {
            return nums;
        }

        int partition = this.partition(nums, 0, nums.length - 1);

        while (partition + 1 != k) {

            if (partition + 1 < k) {
                partition = this.partition(nums, partition + 1, nums.length - 1);
            } else {
                partition = this.partition(nums, 0, partition - 1);
            }
        }

        return Arrays.copyOfRange(nums, 0, partition + 1);

    }

    private final int partition(final int[] nums, final int start, final int end) {

        if (start > end) {
            return -1;
        }

        int pre = start, cur = start;
        final int pivot = nums[end];

        while (cur <= end - 1) {
            if (nums[cur] <= pivot) {
                this.swap(nums, pre, cur);
                ++pre;
            }
            ++cur;
        }

        this.swap(nums, pre, end);
        return pre;

    }

    /**
     *
     */
    private final void swap(final int[] nums, final int a, final int b) {

        final int tmp = nums[a];
        nums[a] = nums[b];
        nums[b] = tmp;

    }

    /**
     * @param arr
     * @param k
     * @return
     */
    private final Integer[] getTopK(final int[] arr, final int k) {

        final PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, Collections.reverseOrder());

        for (int i = 0; i < k; ++i) {
            maxHeap.offer(arr[i]);
        }

        for (int i = k; i < arr.length; ++i) {
            if (arr[i] < maxHeap.peek()) {
                maxHeap.poll();
                maxHeap.offer(arr[i]);
            }
        }
        return maxHeap.toArray(new Integer[0]);

    }
}